import pandas as pd
import os,csv
import argparse


def get_pseudouridine_position_coverage(refPosition, coverage, CYP21A1P_Cov_dict):
    RefPosition = refPosition - 32037580
    QueryStart = 32004844 
    if 398 > RefPosition >= 0 :  # sam位置转换表 IDMS
        queryP = RefPosition
    elif 647 >= RefPosition >= 398 :
        queryP = RefPosition + 1 + 0 - 0
    elif 1276 >= RefPosition > 647 :
        queryP = RefPosition + 2 + 0 - 0
    elif 1369 >= RefPosition > 1276 :
        queryP = RefPosition + 5 + 0 - 0
    elif 1383 >= RefPosition > 1369 :
        queryP = RefPosition + 6 + 0 - 0
    elif 1405 >= RefPosition > 1383 :
        queryP = RefPosition + 9 + 0 - 0
    elif 1420 >= RefPosition > 1405 :
        queryP = RefPosition + 9 + 0 - 1
    elif 1428 >= RefPosition > 1420 :
        queryP = RefPosition + 9 + 0 - 3
    elif 1561 >= RefPosition > 1428 :
        queryP = RefPosition + 11 + 0 - 3
    elif 2603 >= RefPosition > 1561 :
        queryP = RefPosition + 11 + 0 - 11
    elif 4836 >= RefPosition > 2603 :
        queryP = RefPosition + 12 + 0 - 11
    elif 6259 >= RefPosition > 4836 :
        queryP = RefPosition + 13 + 0 - 11
    elif 8548 >= RefPosition > 6259 :
        queryP = RefPosition + 13 + 0 - 131
    elif RefPosition > 8548 :
        queryP = RefPosition + 13 + 375 - 131
    CYP21A1P_coverage = CYP21A1P_Cov_dict[queryP + QueryStart] if (queryP + QueryStart) in CYP21A1P_Cov_dict.keys() else 0
    ratio = 0 if CYP21A1P_coverage == 0 else round(coverage/CYP21A1P_coverage, 5)  # CYP21A1P是0的tatio也是0
    return int(queryP + QueryStart), CYP21A1P_coverage, ratio


def add_cyp21a2_cyp21a1p(true_gene_path, pseudogene_path):
    # true pseudogene cov
    CYP21A2_df=pd.read_csv(true_gene_path, quoting=csv.QUOTE_NONE, sep='\t', names=["chr", "start", "coverage"])
    CYP21A1P_df=pd.read_csv(pseudogene_path, quoting=csv.QUOTE_NONE, sep='\t', names=["chr", "start", "coverage"])
    CYP21A1P_Cov_dict = CYP21A1P_df.set_index('start').coverage.to_dict()
    CYP21A2_df[["CYP21A1P_position", 'CYP21A1P_coverage', 'ratio']] = CYP21A2_df.apply(lambda x: pd.Series(get_pseudouridine_position_coverage(x['start'], x['coverage'], CYP21A1P_Cov_dict)), axis=1)
    return CYP21A2_df


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Get normal sample true and pseudogene cov into the file. Example: add_CYP21A2_CYP21A1P_cov.py -T ${bamSavePath}/${name}_CYP21A2.depth.tsv -P ${bamSavePath}/${name}_CYP21A1P.depth.tsv")
    parser.add_argument('-T', "--true_gene", help="Trun gene cov path.", default="./")
    parser.add_argument('-P', "--pseudogene", help="Pseudogene cov path.", default="./")
    # parser.add_argument('-n', "--number", type=int, help="The number of generate sequence.", default="10")
    args = parser.parse_args()
    resultDf = add_cyp21a2_cyp21a1p(args.true_gene, args.pseudogene)
    resultDf.to_csv("normal_standard_TPCov" + ".csv")
